use riir::agent::prompt::{self, Prompt};
use riir::agent::tools::{ai_task_request, capture_json_code, decode_llm_output};
use riir::llm::Message;
use riir::llm::LLM;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::fmt::Debug;
use std::fs;
use std::sync::Arc;
use tokio::sync::Mutex;
use tokio::io::{self, AsyncWriteExt, Result as IoResult};
use tracing::{debug, info};
use regex::Regex;
use std::path::PathBuf;

pub mod code_trans_agent;
pub mod code_audit_agent;
pub mod code_docs_agent;
pub mod code_refactor_agent;

pub mod code2prompt;

use std::marker::PhantomData;

pub type DynAgentFns = Box<dyn AgentFns>;

#[async_trait]
pub trait AgentFns: Debug + Send {
    // Used to that admin can get attributes from Agents
    fn get_attributes_from_agent(&self) -> &AgentBase;

    // This function will allow agents to execute their logic
    async fn execute(
        &mut self,
        factsheet: &mut FactSheet,
        output_path: PathBuf,
    ) -> Result<(), Box<dyn std::error::Error>>;
}

// // Save Agent Output
// pub fn save_code(contents: &str, path: &str) {
//     let path: String = String::from(path);
//     // 使用正则表达式匹配除了换行符(\n)之外的所有控制字符
//     let control_chars_regex = Regex::new(r"[\x00-\x09\x0B-\x1F\x7F]").unwrap();
//     let cleaned_input = control_chars_regex.replace_all(contents, "").to_string();
//     info!("Save File Content======= {:?}", cleaned_input);
//     fs::write(path, cleaned_input).expect("Failed to write audit file");
// }

async fn save_code(output_file_path: PathBuf, file_suffix: &str, file_content: &str) -> IoResult<()> {
    // 使用正则表达式匹配除了换行符(\n)之外的所有控制字符
    let control_chars_regex = Regex::new(r"[\x00-\x09\x0B-\x1F\x7F]").unwrap();
    let cleaned_content = control_chars_regex.replace_all(file_content, "").to_string();

    let final_path = if output_file_path.is_dir() {
        output_file_path.join(format!("{file_suffix}.md"))
    } else {
        // 如果是文件，为其文件名增加 `_refactor` 后缀
        let file_name = output_file_path.file_stem().unwrap().to_str().unwrap();
        let new_file_name = format!("{}_{}.md", file_name, file_suffix);
        output_file_path.with_file_name(new_file_name)
    };

    // 确保最终路径的父目录存在
    if let Some(parent) = final_path.parent() {
        fs::create_dir_all(parent)?;
    }

    // 写入内容到最终路径指定的文件
    fs::write(final_path, file_content)
}



#[derive(Debug)]
pub struct AgentAdmin {
    attributes: AgentBase,
    agents: Vec<DynAgentFns>,
    // Agent 之间共享消息
    factsheet: FactSheet,
    llm: Arc<Mutex<Box<dyn LLM + 'static>>>,
    output_path: PathBuf,
}

#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
pub struct FactSheet {
    // LLM Output
    pub code: Option<String>,
}

impl AgentAdmin {
    pub async fn new(
        usr_req: String,
        llm: Arc<Mutex<Box<dyn LLM + 'static>>>,
        output_path: PathBuf,
    ) -> Result<Self, Box<dyn std::error::Error>> {
        let position: String = "Agent Admin".to_string();

        let attributes: AgentBase = AgentBase {
            objective: "Rust Lang expert".to_string(),
            position: position.clone(),
            state: AgentState::Starting,
            memory: vec![],
        };

        let factsheet: FactSheet = FactSheet {
            code: Some(usr_req),
        };

        info!("--- factsheet init :  {factsheet:?} ---");

        let agents: Vec<DynAgentFns> = vec![];

        Ok(Self {
            attributes: attributes,
            factsheet,
            agents,
            output_path,
            llm,
        })
    }

    fn add_agent(&mut self, agent: DynAgentFns) {
        self.agents.push(agent);
    }

    fn create_agents(&mut self) {
        self.add_agent(Box::new(code_trans_agent::TransAgent::new(
            self.llm.clone(),
        )));
        self.add_agent(Box::new(code_refactor_agent::RefactorAgent::new(
            self.llm.clone(),
        )));
        self.add_agent(Box::new(code_audit_agent::AuditAgent::new(
            self.llm.clone(),
        )));
        self.add_agent(Box::new(code_docs_agent::DocsAgent::new(self.llm.clone())));
    }

    pub async fn execute(&mut self) {
        self.create_agents();
        info!("Starting Agent Admin! ");

        for agent in &mut self.agents {
            info!("--- EXEC Augent factsheet  :  {:?} ---", self.factsheet);
            let _agent_res: Result<(), Box<dyn std::error::Error>> =
                agent.execute(&mut self.factsheet, self.output_path.clone()).await;
        }
    }
}

pub trait Base {
    fn new(objective: String, position: String) -> Self;
    fn update_state(&mut self, new_state: AgentState);
    fn objective(&self) -> &String;
    fn position(&self) -> &String;
    fn state(&self) -> &AgentState;
    fn memory(&self) -> &Vec<Message>;
}

#[derive(Debug, PartialEq)]
pub enum AgentState {
    Starting,
    Working,
    Debug,
    Done,
}

#[derive(Debug)]
pub struct AgentBase {
    pub objective: String,
    pub position: String,
    pub state: AgentState,
    pub memory: Vec<Message>,
}

impl Base for AgentBase {
    fn new(objective: String, position: String) -> Self {
        Self {
            objective,
            position,
            state: AgentState::Starting,
            memory: Vec::from([]),
        }
    }

    fn update_state(&mut self, new_state: AgentState) {
        self.state = new_state;
    }

    fn objective(&self) -> &String {
        &self.objective
    }

    fn position(&self) -> &String {
        &self.position
    }

    fn state(&self) -> &AgentState {
        &self.state
    }

    fn memory(&self) -> &Vec<Message> {
        &self.memory
    }
}

#[derive(Serialize, Deserialize)]
pub struct RustFixedExpertInfo {
    pub task: String,
}

pub const RUST_FIXED_EXPERT: Prompt<RustFixedExpertInfo> = Prompt(
    r#"

You are an advanced AI assistant with expertise in the Rust programming language. With your advanced machine learning capabilities and comprehensive understanding of Rust, you are tasked with autonomously analyzing Rust code snippets provided by users. Your primary objective is to identify and rectify errors or inefficiencies based on the compilation error messages supplied alongside the code. Your goal is to assist users in debugging and optimizing their Rust code effectively.

**Input Format Provided by Users:**
- `CODE: {:?} \n ERROR_BUGS: {:?} \n`
    - The `CODE` section contains the original Rust code that needs corrections.
    - `ERROR_BUGS` includes the compilation error messages associated with the provided code.

**Your Task:**
- Based on the provided `ERROR_BUGS` messages, analyze the `CODE` to identify the underlying issues.
- Please think step by step, solving errors one at a time, and think carefully through each solution. This careful consideration ensures that each solution is not only effective but also contributes to the overall improvement of the code's quality and maintainability.
- Implement necessary corrections to address these issues, ensuring the code compiles correctly without the previously indicated errors.
- Output the corrected version of the code to the user, providing a direct solution to their programming challenge.

**Operational Guidelines:**
- Thoroughly analyze both the code and the error messages to ensure a comprehensive understanding of the issues at hand.
- Prioritize corrections that directly address the compilation errors, focusing on syntax errors, type mismatches, ownership and borrowing issues, and other common Rust-specific problems.
- Ensure the corrected code is not only free from the identified errors but also adheres to Rust's best practices for readability and efficiency.


**Expected Output:**
Respond in this JSON format:
```json
{{
	"plan": "<concise plan on how you will complete the task>",
	"code": "
	```rust
	<The Rust code after your refactoring,and the Rust code snippet without any escape characters Please!>
	```
	"
}}
```
"#,
    PhantomData,
);
